// Copyright 2013 Dave Pederson.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/**
 * GOSPL - GO web Services PLatform
 *
 * Consolidate and expose HTTP web-services (legacy and new) through a
 * consistent, secure platform. This is accoplished by protecting back-end
 * web-services with tokens. Clients must generate and sign tokens before
 * each request. Tokens will expire after a configurable amount of time.
 */
package main

import (
	"code.google.com/p/go-sqlite/go1/sqlite3"
	"crypto/hmac"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"regexp"
	"strings"
	"time"
)

/**
 * A regex that describes token URL structure
 */
const gosplTokenRegexp = "/[a-f0-9]{40}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{14}"

/**
 * Represents a client request for a service.
 */
type Token struct {
	Signature      string // HMAC-SHA1 signature
	ClientId       string // The Id of the client making the request
	ServiceId      string // The Id of the requested service
	ServiceVersion string // The version of the requested service
	Timestamp      string // The date/time of the request
}

/**
 * Dispatcher for proxying requests to back-end services.
 */
type Dispatcher struct {
	Transport *http.Transport // Needed for proxying requests
}

/**
 * Represents a back-end service.
 */
type Service struct {
	Scheme string // http or https (default http)
	Host   string // IP Address
	Port   int64  // Port (optional)
	URI    string // Replace the token with this value (optional)
}

/**
 * Represents a network connection with a timeout.
 */
type DialTimeoutConn struct {
	net.Conn
	Timeout time.Duration
}

/**
 * Start a web-server and dispatch requests to back-end services.
 */
func main() {
	bindAddress := flag.String(
		"bind", "localhost:80", "The address and port to bind to")
	useSSL := flag.Bool("ssl", false, "Listen / Serve using TLS")
	doThrottle := flag.Bool("throttle", false, "Rate limit the server")
	reqPerSec := flag.Float64("rate", 10.0, "Rate limit req/sec")
	timeout := flag.Int("timeout", 2, "Back-end service timeout")
	flag.Parse()
	log.Println("Starting gospl on", *bindAddress)
	transport := NewDialTimeoutTransport(time.Duration(*timeout) * time.Second)
	dispatcher := Dispatcher{Transport: transport}
	mux := http.NewServeMux()
	if *doThrottle {
		throttle := time.Tick(time.Second / time.Duration(*reqPerSec))
		mux.Handle("/", dispatcher.MakeHandler(throttle))
	} else {
		mux.Handle("/", dispatcher.MakeHandler(nil))
	}
	srv := &http.Server{Handler: mux, Addr: *bindAddress}
	var err error
	if *useSSL {
		err = srv.ListenAndServeTLS("ssl/cert.pem", "ssl/key.pem")
	} else {
		err = srv.ListenAndServe()
	}
	if err != nil {
		log.Fatal("Unable to start web server: ", err)
	}
}

/**
 * Create a http.Transport with read/write timeouts.
 */
func NewDialTimeoutTransport(timeout time.Duration) *http.Transport {
	dialer := func(netw, addr string) (net.Conn, error) {
		c, err := net.DialTimeout(netw, addr, timeout)
		if err != nil {
			return nil, err
		}
		return &DialTimeoutConn{c, timeout}, nil
	}
	return &http.Transport{
		Dial: dialer, DisableKeepAlives: false, DisableCompression: false}
}

/**
 * Read from a connection with a timeout.
 */
func (c *DialTimeoutConn) Read(b []byte) (n int, err error) {
	c.Conn.SetReadDeadline(time.Now().Add(c.Timeout))
	return c.Conn.Read(b)
}

/**
 * Write to a connection with a timeout.
 */
func (c *DialTimeoutConn) Write(b []byte) (n int, err error) {
	c.Conn.SetWriteDeadline(time.Now().Add(c.Timeout))
	return c.Conn.Write(b)
}

/**
 * Hook Dispatcher functionality into a request handler.
 */
func (d *Dispatcher) MakeHandler(throttle <-chan time.Time) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if throttle != nil {
			<-throttle // wait for next 'tick'
		}
		token, err1 := parseToken(r.URL.Path)
		if err1 != nil {
			http.Error(w, err1.Error(), http.StatusBadRequest)
			return
		}
		log.Println("Token:", token)
		conn, err2 := sqlite3.Open("db/gospl.db")
		if err2 != nil {
			http.Error(w, "Internal Error", http.StatusInternalServerError)
			return
		}
		ok, err3 := token.isValid(r.Method, conn)
		if !ok || err3 != nil {
			if err3 != nil {
				http.Error(w, err3.Error(), http.StatusBadRequest)
			} else {
				http.Error(w, "Invalid Token", http.StatusBadRequest)
			}
			return
		}
		service, err4 := token.loadService(conn)
		if err4 != nil {
			http.Error(w, err4.Error(), http.StatusBadRequest)
			return
		}
		d.DispatchRequest(w, r, service)
	}
}

/**
 * Convert a URL string into a token.
 */
func parseToken(url string) (token *Token, err error) {
	var ok bool
	ok, err = regexp.Match(gosplTokenRegexp, []byte(url))
	if !ok || err != nil {
		if err != nil {
			log.Println(err)
		}
		err = errors.New("Token not found in URL")
		return
	}
	fields := strings.Split(url, "/")
	tokenString := fields[1]
	tokenFields := strings.Split(tokenString, "-")
	token = &Token{
		Signature:      tokenFields[0],
		ClientId:       tokenFields[1],
		ServiceId:      tokenFields[2],
		ServiceVersion: tokenFields[3],
		Timestamp:      tokenFields[4],
	}
	return
}

/**
 * Load the appropriate service for a given token.
 */
func (token *Token) loadService(conn *sqlite3.Conn) (service *Service, err error) {
	args := sqlite3.NamedArgs{"$service_id": token.ServiceId, "$service_version": token.ServiceVersion}
	sql := "SELECT * FROM services WHERE service_id = $service_id AND service_version = $service_version"
	s, err := conn.Query(sql, args)
	if err == nil {
		row := make(sqlite3.RowMap)
		s.Scan(row)
		service = &Service{
			Scheme: row["scheme"].(string),
			Host:   row["host"].(string),
			Port:   row["port"].(int64),
			URI:    row["endpoint"].(string),
		}
	}
	return
}

/**
 * Determine if a token is valid.
 */
func (token *Token) isValid(method string, conn *sqlite3.Conn) (ok bool, err error) {
	var secKey string
	secKey, err = token.validateClient(conn)
	if err != nil {
		log.Println(err)
		err = errors.New("Invalid client")
		return
	}
	var ttl int64
	ttl, err = token.validateClientPermission(conn)
	if err != nil {
		log.Println(err)
		err = errors.New("Permission denied")
		return
	}
	ok, err = token.validateSignature(secKey, method)
	if !ok || err != nil {
		if err != nil {
			log.Println(err)
		}
		err = errors.New("Invalid signature")
		return
	}
	ok, err = token.validateTimeout(ttl)
	if !ok || err != nil {
		if err != nil {
			log.Println(err)
		}
		err = errors.New("Token has expired")
		return
	}
	ok = true
	return
}

/**
 * Determine if client exists.  If so, return the client secret key.
 */
func (token *Token) validateClient(conn *sqlite3.Conn) (secKey string, err error) {
	args := sqlite3.NamedArgs{"$client_id": token.ClientId}
	sql := "SELECT * FROM clients WHERE client_id = $client_id"
	s, err := conn.Query(sql, args)
	if err == nil {
		row := make(sqlite3.RowMap)
		s.Scan(row)
		secKey = row["secret_key"].(string)
	}
	return
}

/**
 * Determine if a client has permission to call a service.
 */
func (token *Token) validateClientPermission(conn *sqlite3.Conn) (ttl int64, err error) {
	args := sqlite3.NamedArgs{"$client_id": token.ClientId, "$service_id": token.ServiceId, "$service_version": token.ServiceVersion}
	sql := "SELECT * FROM permissions WHERE client_id = $client_id AND service_id = $service_id AND service_version = $service_version"
	s, err := conn.Query(sql, args)
	if err == nil {
		row := make(sqlite3.RowMap)
		s.Scan(row)
		ttl = row["ttl"].(int64)
	}
	return
}

/**
 * Determine if a token signature is valid.
 */
func (token *Token) validateSignature(secKey, method string) (bool, error) {
	// Calculate the HMAC SHA1 hash of a data string.
	hmacSha1 := func(data, key string) []byte {
		hash := hmac.New(sha1.New, []byte(key))
		hash.Write([]byte(data))
		return hash.Sum(nil)
	}
	data := method + "-" + token.ClientId + "-" + token.ServiceId + "-" +
		token.ServiceVersion + "-" + token.Timestamp
	expected := hmacSha1(data, secKey)
	provided, err := hex.DecodeString(token.Signature)
	if err != nil {
		return false, err
	}
	if subtle.ConstantTimeCompare(expected, provided) != 1 {
		log.Println("Expected:", hex.EncodeToString(expected))
		log.Println("Provided:", token.Signature)
		return false, errors.New("Invalid signature")
	}
	return true, nil
}

/**
 * Determine if a token has passed a time-to-live expiration.
 */
func (token *Token) validateTimeout(ttl int64) (ok bool, err error) {
	var age int64
	age, err = token.Age()
	if err != nil {
		ok = false
		return
	}
	log.Println("Token age:", age, ", ttl:", ttl)
	if age <= ttl {
		ok = true
	} else {
		err = errors.New("Expired token")
	}
	return
}

/**
 * Determine the age of a token
 */
func (token *Token) Age() (int64, error) {
	ts, err := time.Parse("20060102150405", token.Timestamp)
	if err != nil {
		return -1, err
	}
	tsNow := time.Now().UTC().Unix()
	tsUnix := ts.Unix()
	if tsNow < tsUnix {
		return -1, errors.New("Invalid future timestamp")
	}
	return tsNow - tsUnix, nil
}

/**
 * Dispatch a request to a back-end service.
 */
func (d *Dispatcher) DispatchRequest(
	w http.ResponseWriter, r *http.Request, service *Service) {

	log.Println("Request URL:", r.URL)
	if service.Scheme != "" {
		r.URL.Scheme = service.Scheme
	} else {
		log.Println("No Scheme: using http")
		r.URL.Scheme = "http"
	}
	setXForwardedFor(r)

	// Set back-end host and port (if provided)
	if service.Port != 0 {
		r.URL.Host = fmt.Sprintf("%s:%d", service.Host, service.Port)
	} else {
		r.URL.Host = service.Host
	}

	// Re-write token using back-end service URI
	rewrite, _ := regexp.Compile(gosplTokenRegexp)
	r.URL.Path = rewrite.ReplaceAllString(r.URL.Path, service.URI)

	// Proxy request
	log.Println("Re-writing request to:", r.URL)
	resp, err := d.Transport.RoundTrip(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()

	// Copy response headers and copy proxy data to response
	for k, v := range resp.Header {
		for _, vv := range v {
			w.Header().Add(k, vv)
		}
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

/**
 * Determine remote address (client IP) and add an X-Forwarded-For header.
 */
func setXForwardedFor(r *http.Request) {
	if _, ok := r.Header["X-Forwarded-For"]; !ok {
		remoteAddr := r.RemoteAddr
		idx := strings.LastIndex(remoteAddr, ":")
		if idx != -1 {
			remoteAddr = remoteAddr[0:idx]
			if remoteAddr[0] == '[' && remoteAddr[len(remoteAddr)-1] == ']' {
				remoteAddr = remoteAddr[1 : len(remoteAddr)-1]
			}
		}
		r.Header.Add("X-Forwarded-For", remoteAddr)
	}
}
